#!/usr/bin/env python
# -*- coding: utf-8 -*-
###############################################################################
#
# A script for expanding alignments, avoiding use of ambiguous symbols in the
# the sequences.
# Developed by Francisco Gabriel Malbrán [fgmalbran@gmail.com]
#
###############################################################################
#
# FuDePAN
# Copyright (C) Francisco Gabriel Malbran, 2009, FuDePAN
#
#   fasta-expander is free software: you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#   fasta-expander is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with fasta-expander.  If not, see <http://www.gnu.org/licenses/>.
#
###############################################################################
import re, getopt, sys, os.path as path, datetime

NON_AMBIGUOS_CHARS = ['A','C','G','T','U','-']
AMBIGUOS_CHARS = ['R','Y','K','M','S','W','B','D','H','V','N']
UNAMBIG_DICT = { 
                'R':['G','A'], 
                'Y':['T','C'],
                'K':['G','T'],
                'M':['A','C'],
                'S':['G','C'],
                'W':['A','T'],
                'B':['G','T','C'],
                'D':['G','A','T'],
                'H':['A','C','T'],
                'V':['G','C','A'],
                'N':['A','G','C','T']
                }
ALLOWED_CHARS = ['A','T','C','G','U', '-', '\n'] # non-ambiguous chars still allowed.
HEADER_FORMAT = ">.+\n"
_debug = True   # Debug mode is default mode.
_count = 0      # Quantity of extensions to make/made
_verbose = False    # Verbose mode
_ranges = []        # for Ranges mode

def main(argv):
    time = datetime.datetime.now()  # for simple time measure
    global _debug, _count, _verbose, AMBIGUOS_CHARS, _ranges
    prev = ''   # First part of the alignment file (before first header, if any)
    head = ''   # The header of the sequence being analized
    seq = ''    # A sequence of the alignment
    filename = ''   # Name of the alignment file
    #-- option handle
    try:
        opts, args = getopt.getopt(argv, "f:hgvi:r:")
    except getopt.GetoptError, error:
        sys.stderr.write(str(error) + "\nGet usage info with -h\n")
        sys.exit()
    for opt, arg in opts:
        if opt == "-h":     # help
            usage()
            sys.exit()
        elif opt == "-g":   # generate
            _debug = False
        elif opt == "-v":   # verbose
            _verbose = True
        elif opt == "-i":   # ignore
            AMBIGUOS_CHARS = [x for x in AMBIGUOS_CHARS if x not in list(arg.upper())]
        elif opt == "-f":   # file
            filename = arg
        elif opt == "-r":   # ranges
            try:
                p = open(arg, 'r')
            except IOError:
                sys.stderr.write("Ranges file unknown or unspecified. Get usage info with -h\n")
                sys.exit()
            for line in p:  # validation of ranges file format + parse of ranges
                # ENHANCE: Check order and overlap of ranges defined!
                m = re.match("\s*(\d+)\s*,\s*(\d+)\s*$", line)
                if m:
                    _ranges.append([int(x) for x in m.groups()])
                # allow white lines and comments with #
                elif not (re.match("^\s*$", line) or re.match("^#.*$",line)):
                    sys.stderr.write("Bad format in ranges file. Get usage info with -h\n")
                    p.close()
                    sys.exit()
            p.close()
    #-- main program
    try:
        f = open(filename, 'r')
    except IOError:
        sys.stderr.write("File unknown or unspecified. Get usage info with -h\n")
        sys.exit()
    # get first header
    while True:
        line = f.readline()
        if line == '': # means EOF reached
            sys.stderr.write("Bad alignment file. Could not find a sequence header.\n")
            sys.exit()
        if re.match(HEADER_FORMAT, line):
            head = line
            break
        prev += line
    # write prev, if needed
    if not _debug:
        sys.stdout.write(prev)
    # proccess/analize by sequences
    while True:
        line = f.readline()
        if re.match(HEADER_FORMAT, line) or line == '':
            if not _debug:
                sys.stdout.write(unambig(head, seq))
            else:
                updateCount(head, seq)
            head = line
            seq = ''
        else:
            seq += line
        if line == '': break    # '' means EOF reached
    # end
    f.close()
    if _debug:
        sys.stderr.write("%.5G extensions would be needed to proccess this file.\n" % _count)
    else:
        sys.stderr.write("%.5G extensions made.\n" % _count)
    sys.stderr.write("Analyzed/processed in %s. (hh.mm.ss)\n\n" % (datetime.datetime.now()-time))
    return

def unambig(head, seq):
# Makes all possible extensions on a sequence
    global _count
    local_count = 1 # Count of extensions for this seq
    result = ''     # Final extended chunk
    new_seq = ''    # Pre-processed seq
    rep_list = []   # List of replacement char lists
    pos = 1         # Position of the char being analized
    for ch in seq:
        ch = ch.upper()
        if ch in AMBIGUOS_CHARS:
            if isAllowed(pos): # if inside a range or not range defined, unambiguate
                local_count *= UNAMBIG_DICT.get(ch).__len__()
                rep_list.append(UNAMBIG_DICT.get(ch))
                new_seq += '%s'
            else: # else add as it is
                new_seq += ch
            pos += 1
        else:
            if ch not in ALLOWED_CHARS:
                sys.stderr.write("** Warning!:Unknown char found: \'%s\' on header: %s**\n" % (ch, head))
            else:
                if ch != '\n':
                    pos += 1
                new_seq += ch
    # "unambiguate" seq if neccesary
    if rep_list:
        _count += local_count
        j = 1
        for t in list(apply(icross,rep_list)):
            sub_head = "%s_SPLITTED_%s\n"%(head.strip(), j)
            result += sub_head + new_seq % t + '\n'
            j += 1
    else:
        result = head + new_seq
    return result

def updateCount(head, seq):
# Only analizes seq updating general extension counter
    global _count, _verbose
    local_count = 1
    for match in re.finditer(str(AMBIGUOS_CHARS), seq, re.IGNORECASE):
        if isAllowed(match.start()+1):
            local_count *= UNAMBIG_DICT.get(match.group().upper()).__len__()
    if local_count > 1:
        _count += local_count
        if _verbose: sys.stderr.write("head: %s --> %.5G exts\n" % (head, local_count))
    return

def icross(*sequences):
# Creates the cartesian product of a list of seqs
# ENHANCE: Maybe put this into a library?
    if sequences:
        for x in sequences[0]:
            for y in icross(*sequences[1:]):
                yield(x,)+y
    else: yield()
    return

def isAllowed(pos):
# Determines if the position "pos" is within specified analisys range (if any)
    if _ranges: # do kind of a binary search
        return bSearchForRanges(_ranges, pos)
    else:
        return True

def bSearchForRanges(seqOfRanges, elm):
# Does a binary-like search to check if elm is within 
# any of the ranges in seqOfRanges
# IMPORTANT: It is a precondition for this to work, that the seqOfRanges object 
# IS ORDERED, meaning that elements of a flatten version of seqOfRanges should 
# be in strict ascendant order.
# ENHANCE: Support unordered sets of lists and detect overlapping.
# ENHANCE: Maybe put into a library?
    left = 0
    right = len(seqOfRanges) - 1
    while left <= right:
        center =  (left + right)/2
        cRan = seqOfRanges[center]
        if cRan[0] <= elm and elm <= cRan[1]:
            return True
        elif elm < cRan[0]:
            right = center - 1
        else:
            left = center + 1
    return False


def usage():
# Prints usage info
    sys.stdout.write("Usage info:\nCall the script as\n\t./fasta-expander.py -f filename [-g] [-r rangesfile]\n"+\
            "where \'filename\' is an alignment file you want to analize/process.\n"+\
            "The script will show how many extensions would be needed to make\n"+\
            "the alignment unambiguos.\n"+\
            "-g: Generate. Show extended file in standar output.\n"+\
            "-v: Verbose. Shows number of extensions needed for each ambiguous\n"+\
                "sequence (does not work together with -g)\n"+\
            "-i XXXX: Ignore. Treats the characters of the following string as\n"+\
                "non-ambiguous. Use example:\n"+\
                "\t./fasta-expander.py -gf alignment_file.txt -i NRK.\n"+\
                "will treat pseudo-nucleotids N, R, and K as regular nucleotids.\n"+\
            "-r rangesfile: Ranges. Helps you decide which exact parts of each\n"+\
                "sequence will be analized/processed through the reading of a\n"+\
                "file with content in the following format:\n"+\
                "\tbegin1,end1\n\tbegin2,end2\n\tetc.\n"+\
                "where the begins and ends are int numbers in crescent order starting at 1.\n"+\
                "Spaces, blank lines, and comments with '#' are allowed.\n"+\
            "-h: Shows this help message\n")
    return

# Main call
if __name__ == "__main__":
    main(sys.argv[1:])
